"""
将拍摄到的视频文件转换为图像序列，用于训练神经网络

date: 2025/9/27
author: SiHeng Tang
file: v2p.py
copyright(c) DFSA Eastern Software
此程序不能提供任何担保 WITHOUT WARRANTY OF ANY KIND
"""
import argparse
import logging
import os

import cv2

logger = logging.getLogger(__name__)
DEFAULT_DIRECTORY_NAME = "picture_sequence"


def extract_frames(video_path, interval, start_time, end_time, output_dir):
    """
    将视频拆分为图像帧
    :param video_path: 视频文件目录
    :param interval: 拆分间隔，秒
    :param start_time: 开始时间，秒
    :param end_time: 结束时间，秒
    :param output_dir: 输出目录
    """
    cap = cv2.VideoCapture(video_path)

    if not cap.isOpened():
        logger.error("Cannot open video file.")
        exit(0)

    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = total_frames / fps

    if start_time is None:
        start_time = 0
    if end_time is None:
        end_time = duration

    if start_time < 0 or end_time > duration or start_time >= end_time:
        logger.error("Invalid time range.")
        exit(0)

    start_frame = int(start_time * fps)
    end_frame = int(end_time * fps)
    frame_interval = int(interval * fps)

    base_name = os.path.splitext(os.path.basename(video_path))[0]

    os.makedirs(output_dir, exist_ok=True)

    logger.info(f"Extracting frames from {start_time}s to {end_time}s, every {interval}s...")

    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
    frame_idx = start_frame
    count = 0

    while frame_idx <= end_frame:
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
        ret, frame = cap.read()

        if not ret:
            break

        output_path = os.path.join(output_dir, f"{base_name}_{count:06d}.jpeg")
        cv2.imwrite(output_path, frame)
        logger.info(f"Saved: {output_path}")

        frame_idx += frame_interval
        count += 1

    cap.release()
    logger.info("Extraction complete.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Extract frames from video at regular intervals.")
    parser.add_argument("video_path", help="Path to the video file")
    parser.add_argument("interval", type=float, help="Interval between frames in seconds")
    parser.add_argument("--output", default=DEFAULT_DIRECTORY_NAME,
                        help=f"Output directory (default: '{DEFAULT_DIRECTORY_NAME}')")
    parser.add_argument("--start", type=float, help="Start time in seconds (default: 0)")
    parser.add_argument("--end", type=float, help="End time in seconds (default: full length)")

    args = parser.parse_args()
    extract_frames(args.video_path, args.interval, args.start, args.end, args.output)
